from modelscope import  snapshot_download
from FlagEmbedding import FlagModel
import base_util
from BCEmbedding import EmbeddingModel


#
class BaseEmbeddingModel:
    """
    抽象基类，定义了Embedding模型的接口。
    """
    def __init__(self, model_name):
        # 这里的 模型下载路径 目前我是默认在当前目录下下载的。 显示场景中， 模型下载路径或者加载路径应该由 配置文件去定义。

        self.model_download_path=f"{base_util.current_project_dir}\\\models\EmbeddingModels\\"
        self.model_name=model_name
        self.model=None

    def embed(self, text):
        """
        将文本转换为嵌入向量。
        :param text: 输入文本
        :return: 嵌入向量
        """
        return self.model.encode(text, convert_to_tensor=True)


    def load_model(self,model_path):
        # load_model的抽象方法。 不同模型有不同的加载方式。 自己扩展
        pass

    def download_model(self,model_id,model_path):
        model_dir = snapshot_download(model_id, cache_dir=model_path)
        return model_dir


    def similarity(self, embeddings1, embeddings2):
        """
        计算两组嵌入向量之间的相似度。
        :param embeddings1: 第一组嵌入向量
        :param embeddings2: 第二组嵌入向量
        :return: 相似度矩阵
        """
        return (embeddings1 @ embeddings2.T).numpy()

class BgeEmbedding(BaseEmbeddingModel):
    """
    使用bge-large-zh-v1.5模型的嵌入类。
    """
    def __init__(self, model_name="AI-ModelScope/bge-large-zh-v1.5"):
        # super().__init__('BAAI/bge-large-zh-v1.5')
        super().__init__(model_name)
        self.load_model(model_name)


    def load_model(self,model_id):
        # 下载模型 直接从 ModelScope 下载。 (被墙的情况下。  后期会增加动态切换功能。即。 先判断当前网络。再决定从哪里下,为了方便 这里直接从modeScope中下载)
        model_dir = snapshot_download(model_id, cache_dir=self.model_download_path)

        model = FlagModel(model_dir,
                          query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章：",
                          use_fp16=True)  # Setting use_fp16 to True speeds up computation with a slight performance degradation
        self.model=model
        return model


    def embed(self, text):
        """
            将文本转换为嵌入向量。
            :param text: 输入文本
            :return: 嵌入向量
            """
        return self.model.encode(text)

    # 示例使用


class BCEmbedding(BaseEmbeddingModel):

    def __init__(self, model_name='maidalun/bce-embedding-base_v1'):
        # super().__init__('BAAI/bge-large-zh-v1.5')
        super().__init__(model_name)
        self.load_model(model_name)

    def load_model(self, model_id='maidalun/bce-embedding-base_v1'):
        model_dir = snapshot_download(model_id, cache_dir=self.model_download_path)
        self.model = EmbeddingModel(model_name_or_path=model_dir)
        return self.model

    def embed(self, text):
        """
            将文本转换为嵌入向量。
            :param text: 输入文本
            :return: 嵌入向量
            """
        return self.model.encode(text)


if __name__ == "__main__":
    # 初始化Embedding模型
    # 调用Embedding模型的简单办法
    embedding_model = BCEmbedding()
    s=embedding_model.embed(["你好"])
    print(s)
